
import argparse
import re
from funasr import AutoModel
from paddlespeech.cli.text.infer import TextExecutor
import os
from pathlib import Path

# 统一模型下载/缓存目录到项目内的 models 目录
MODELS_DIR = Path(os.path.dirname(__file__)) / "models"
MODELS_DIR.mkdir(exist_ok=True)
os.environ["MODELSCOPE_CACHE"] = MODELS_DIR.as_posix()
os.environ["HUGGINGFACE_HUB_CACHE"] = MODELS_DIR.as_posix()

def transcribe_audio(audio_path):
    """
    Transcribes an audio file using the SenseVoice model.
    """
    model = AutoModel(
        model="iic/SenseVoiceSmall",
        vad_model="fsmn-vad",
        vad_kwargs={"max_single_segment_time": 30000},
        device="cpu",
        cache_dir=MODELS_DIR.as_posix(),
    )
    res = model.generate(
        input=audio_path,
        cache={},
        language="auto",
        use_itn=True,
        batch_size_s=60,
        merge_vad=True,
        merge_length_s=15
    )
    return res[0]["text"]

def punctuate_text(text):
    """
    Adds punctuation to a text using the PaddleSpeech model.
    """
    text_punc = TextExecutor()
    result = text_punc(text=text)
    return result

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Speech recognition with punctuation.")
    parser.add_argument("audio_file", type=str, help="Path to the audio file.")
    parser.add_argument("--output_file", type=str, help="Optional path to save the output.")
    args = parser.parse_args()

    transcribed_text = transcribe_audio(args.audio_file)
    # 打印转录文本
    print("Transcription:", transcribed_text)
    # Remove special tokens from the transcription
    transcribed_text = re.sub(r'<\|.*?\|>', '', transcribed_text)
    # punctuated_text = punctuate_text(transcribed_text)

    if args.output_file:
        with open(args.output_file, "w", encoding="utf-8") as f:
            f.write("Transcription:\n")
            f.write(transcribed_text + "\n\n")
            # f.write("Punctuated text:\n")
            # f.write(punctuated_text + "\n")
    else:
        print("Transcription:")
        print(transcribed_text)
        # print("\nPunctuated text:")
        # print(punctuated_text)
